from re import A
import requests
import numpy as np

class LLMProxy:
    def __init__(self, embedding_api: str, embedding_key: str, embedding_name: str, reranker_api: str, reranker_key: str, reranker_name: str):
        self.embedding_api = embedding_api
        self.embedding_key = embedding_key
        self.embedding_name = embedding_name
        self.reranker_api = reranker_api
        self.reranker_key = reranker_key
        self.reranker_name = reranker_name
    
    def embedding(self, query):
        payload = {
            "model": self.embedding_name,
            "input": query,
            "encoding_format": "float"
        }
        headers = {
            "Authorization": f"Bearer {self.embedding_key}",
        }

        for _ in range(3):
            try:
                response = requests.post(f"{self.embedding_api}/embeddings", json=payload, headers=headers)
                if response.status_code != 200:
                    raise response.text
                resdata = response.json()
                return np.array([item["embedding"] for item in resdata["data"]]).astype(np.float32)
            except Exception as e:
                print(str(e))
    
    def rerank_score(self, query, documents):

        payload = {
            "model": self.reranker_name,
            "query": query,
            "documents": documents,
            "return_documents": False,
        }
        headers = {
            "Authorization": f"Bearer {self.reranker_key}",
            "Content-Type": "application/json"
        }

        for _ in range(3):
            try:
                response = requests.request("POST", f"{self.reranker_api}/rerank", json=payload, headers=headers)

                if response.status_code != 200:
                    raise response.text
                
                resdata = response.json()
                
                return [item["relevance_score"] for item in resdata["results"]]
            except Exception as e:
                print(str(e))